import cv2
import numpy as np
import os
import os.path as osp

IMG_SIZE = 3000

DATADIR = os.path.abspath(".")
IMAGEDIR = os.path.join(DATADIR, 'val_img')
LABELDIR = os.path.join(DATADIR, 'val_txt')

OUTDIR = os.path.join(DATADIR, 'val_watch')
if not os.path.exists(OUTDIR):
    os.makedirs(OUTDIR)


tif_names = set([name[:-4] for name in os.listdir(IMAGEDIR) if name.endswith('.tif')])
txt_names = set([name[:-4] for name in os.listdir(LABELDIR) if name.endswith('.txt')])
intersection_names = tif_names & txt_names


names = os.listdir(IMAGEDIR)
for name in names:
    name = name.split(".")[0]
    
    if name not in intersection_names:
        continue

    image = cv2.imread(osp.join(IMAGEDIR, name+".tif"))    
    with open(osp.join(LABELDIR, name+".txt"), 'r') as fp:
        lines = fp.readlines()

    for line in lines:
        line = line.split(" ")
        label = float(line[0])
        cx = IMG_SIZE * float(line[1])
        cy = IMG_SIZE * float(line[2])
        w = IMG_SIZE * float(line[3])
        h = IMG_SIZE * float(line[4])
        angle = int(line[5])
        arg = ((cx, cy), (w, h), (angle))
        box = cv2.boxPoints(arg)
        box = np.int0(box)
        cv2.drawContours(image, [box], 0, (0, 255, 255), 5)
        # cv2.imshow(" ", image)

    cv2.imwrite(osp.join(OUTDIR, name+".tif"), image)
    print(name + ", done...")
